# coding: utf-8

from numpy import array, sqrt, inf


def find_correspondence(target_cloud: array, weight=(1, 1, 1)):
    """ Find correspondences in a plain way.
    :param target_cloud: the target cloud.
    :param weight: the weight.
    :return: the array of correspondences.
    """

    def correspondence(p):
        dist = pow((target_cloud - p) * weight, 2).sum(axis=1)
        min_idx = dist.argmin()
        return p, target_cloud[min_idx, :], min_idx

    return correspondence


def reject_correspondence(max_dist=50, max_y=30):
    def reject(correspondence):
        p, x = correspondence[:2]
        dist = sqrt(pow(p - x, 2).sum())
        return abs(p[1] - x[1]) < max_y and dist < max_dist

    return reject


def find_correspondences(source_cloud, target_cloud, max_dist=inf):
    source_correspondences = list(map(
            # find closest point.
            find_correspondence(target_cloud),
            source_cloud
    ))

    filtered_correspondences = list(filter(reject_correspondence(max_dist), source_correspondences))

    print(">>> correspondences: %d" % len(filtered_correspondences))

    return [array(p) for p in zip(*map(
        lambda c: c[:2], filtered_correspondences
    ))]